from time import sleep
import json
import re
import os
from datetime import datetime
import requests

class LangchainChatWrapper:
    def __init__(self, args):
        """
        初始化聊天包装器

        :param args: 包含所有配置参数的对象
        """
        self.args = args
        self.openai_key = args.openai_key
        self.model_name = args.model_name
        self.openai_base = args.openai_base
        self.temperature = args.temperature
        self.max_length = args.max_length
        self.data_path = args.data_path
        self.output_path = args.output_path
        self.generation_strategy = args.generation_strategy  # 0: 整体, 1: 迭代
        self.sample = args.sample
        
        # 对话历史记录
        self.conversation_history = []
        
        # API配置将在请求时使用
        
        # 加载数据源
        self.problems = self.load_problems()
        
    def load_problems(self):
        """
        从JSON文件加载问题数据
        
        :return: 问题列表
        """
        try:
            with open(self.data_path, 'r', encoding='utf-8') as f:
                problems = json.load(f)
            print(f"成功加载 {len(problems)} 个问题")
            return problems
        except Exception as e:
            print(f"加载数据失败: {str(e)}")
            return []
    
    def extract_python_code(self, text):
        """
        从markdown格式的文本中提取Python代码
        
        :param text: 包含代码的文本
        :return: 提取的Python代码
        """
        # 匹配```python 或 ``` 代码块
        patterns = [
            r'```python\s*\n(.*?)\n```',
            r'```\s*\n(.*?)\n```',
            r'```python(.*?)```',
            r'```(.*?)```'
        ]
        
        for pattern in patterns:
            matches = re.findall(pattern, text, re.DOTALL)
            if matches:
                # 返回第一个匹配的代码块
                code = matches[0].strip()
                if code:
                    return code
        
        # 如果没有找到代码块，尝试提取看起来像Python代码的部分
        lines = text.split('\n')
        code_lines = []
        in_code = False
        
        for line in lines:
            # 检查是否是Python代码行
            if (line.strip().startswith(('def ', 'class ', 'import ', 'from ', 'if ', 'for ', 'while ', 'try:', 'except', 'return')) or
                ('=' in line and not line.strip().startswith('#')) or
                line.strip().startswith(('    ', '\t'))):
                in_code = True
                code_lines.append(line)
            elif in_code and line.strip() == '':
                code_lines.append(line)
            elif in_code and not line.strip().startswith('#'):
                code_lines.append(line)
            elif in_code:
                break
        
        if code_lines:
            return '\n'.join(code_lines).strip()
        
        return text.strip()
    
    def create_prompt(self, problem, iteration_feedback=None):
        """
        为问题创建提示词
        
        :param problem: 问题数据
        :param iteration_feedback: 迭代反馈（用于迭代策略）
        :return: 提示词
        """
        base_prompt = f"""
请解决以下编程问题：

问题描述：
{problem['question']}

输入输出示例：
{problem['input_output']}

请提供完整的Python解决方案。代码应该能够处理所有给定的测试用例。
"""
        
        if iteration_feedback:
            base_prompt += f"""

之前的尝试存在以下问题：
{iteration_feedback}

请根据反馈改进代码。
"""
        
        return base_prompt
    
    def create_feedback_prompt(self, problem, code, test_results=None):
        """
        创建反馈提示词（用于迭代策略）
        
        :param problem: 问题数据
        :param code: 生成的代码
        :param test_results: 测试结果
        :return: 反馈提示词
        """
        feedback_prompt = f"""
请分析以下代码是否正确解决了问题：

问题：
{problem['question']}

代码：
{code}

输入输出要求：
{problem['input_output']}

请指出代码中可能存在的问题，包括：
1. 逻辑错误
2. 边界条件处理
3. 输入输出格式
4. 算法效率

请提供具体的改进建议。
"""
        
        if test_results:
            feedback_prompt += f"""

测试结果：
{test_results}
"""
        
        return feedback_prompt
    
    def generate_result(self, _prompt, reset_context=False, retry_count=0, max_retries=3):
        """
        根据给定的提示信息生成结果，可选择是否重置对话上下文

        :param _prompt: 输入的提示信息
        :param reset_context: 是否重置对话上下文（默认False）
        :param retry_count: 当前重试次数
        :param max_retries: 最大重试次数
        :return: 生成的回复内容
        """
        if reset_context:
            self.conversation_history = []  # 清空对话历史

        try:
            # 添加用户消息到历史
            self.conversation_history.append({"role": "user", "content": _prompt})

            # 准备API请求
            headers = {
                "Authorization": f"Bearer {self.openai_key}",
                "Content-Type": "application/json"
            }
            
            data = {
                "model": self.model_name,
                "messages": self.conversation_history,
                "temperature": self.temperature,
                "max_tokens": self.max_length
            }
            
            # 发送请求
            response = requests.post(
                f"{self.openai_base}/chat/completions",
                headers=headers,
                json=data,
                timeout=60
            )
            
            if response.status_code == 200:
                result = response.json()
                ai_message = result['choices'][0]['message']['content']
                
                # 添加AI回复到历史
                self.conversation_history.append({"role": "assistant", "content": ai_message})
                
                return ai_message
            else:
                raise Exception(f"API请求失败: {response.status_code}, {response.text}")

        except Exception as e:
            if retry_count < max_retries:
                print(f"generation failure (重试 {retry_count + 1}/{max_retries}): {str(e)}")
                sleep(2)  # 减少等待时间
                return self.generate_result(_prompt, reset_context, retry_count + 1, max_retries)
            else:
                print(f"generation failure (已达到最大重试次数): {str(e)}")
                # 返回一个模拟的响应，避免程序崩溃
                mock_response = f"# 模拟响应 - API调用失败\n# 原始提示词长度: {len(_prompt)} 字符\n\ndef mock_solution():\n    # 这是一个模拟的解决方案\n    # 由于API调用失败，无法生成真实代码\n    pass\n\nprint('API调用失败，这是模拟响应')"
                
                # 添加模拟回复到历史
                self.conversation_history.append({"role": "assistant", "content": mock_response})
                
                return mock_response
    
    def generate_holistic(self, problem):
        """
        整体生成策略：为单个问题生成多个独立的解决方案
        
        :param problem: 问题数据
        :return: 生成结果列表
        """
        results = []
        prompt = self.create_prompt(problem)
        
        for i in range(self.sample):
            print(f"  生成第 {i+1}/{self.sample} 个解决方案...")
            
            # 每次生成都重置上下文，确保独立性
            response = self.generate_result(prompt, reset_context=True)
            code = self.extract_python_code(response)
            
            result = {
                'iteration': i + 1,
                'prompt': prompt,
                'raw_response': response,
                'extracted_code': code,
                'timestamp': datetime.now().isoformat()
            }
            results.append(result)
            
            sleep(1)  # 避免API调用过于频繁
        
        return results
    
    def generate_incremental(self, problem):
        """
        迭代生成策略：通过反馈不断改进代码
        
        :param problem: 问题数据
        :return: 生成结果列表
        """
        results = []
        current_code = None
        feedback = None
        
        for i in range(self.sample):
            print(f"  迭代第 {i+1}/{self.sample} 轮...")
            
            # 创建提示词
            prompt = self.create_prompt(problem, feedback)
            
            # 生成代码
            if i == 0:
                response = self.generate_result(prompt, reset_context=True)
            else:
                response = self.generate_result(prompt, reset_context=False)
            
            current_code = self.extract_python_code(response)
            
            result = {
                'iteration': i + 1,
                'prompt': prompt,
                'raw_response': response,
                'extracted_code': current_code,
                'feedback_used': feedback,
                'timestamp': datetime.now().isoformat()
            }
            results.append(result)
            
            # 如果不是最后一轮，生成反馈
            if i < self.sample - 1:
                feedback_prompt = self.create_feedback_prompt(problem, current_code)
                feedback_response = self.generate_result(feedback_prompt, reset_context=False)
                feedback = feedback_response
                
                sleep(1)  # 避免API调用过于频繁
        
        return results
    
    def run_generation(self):
        """
        运行代码生成流程
        
        :return: 所有结果
        """
        all_results = []
        
        print(f"开始生成，共 {len(self.problems)} 个问题")
        print(f"生成策略: {'迭代' if self.generation_strategy == 1 else '整体'}")
        print(f"每个问题生成 {self.sample} 次")
        
        for idx, problem in enumerate(self.problems):
            print(f"\n处理问题 {idx+1}/{len(self.problems)} (ID: {problem['id']})")
            
            try:
                if self.generation_strategy == 1:  # 迭代策略
                    results = self.generate_incremental(problem)
                else:  # 整体策略
                    results = self.generate_holistic(problem)
                
                problem_result = {
                    'problem_id': problem['id'],
                    'problem_url': problem['url'],
                    'difficulty': problem['difficulty'],
                    'generation_strategy': 'incremental' if self.generation_strategy == 1 else 'holistic',
                    'results': results,
                    'total_generations': len(results)
                }
                
                all_results.append(problem_result)
                
            except Exception as e:
                print(f"处理问题 {problem['id']} 时出错: {str(e)}")
                continue
        
        # 保存结果
        self.save_results(all_results)
        
        return all_results
    
    def save_results(self, results):
        """
        保存生成结果到文件
        
        :param results: 生成结果
        """
        try:
            output_data = {
                'generation_config': {
                    'model_name': self.model_name,
                    'temperature': self.temperature,
                    'max_length': self.max_length,
                    'generation_strategy': 'incremental' if self.generation_strategy == 1 else 'holistic',
                    'sample_count': self.sample,
                    'data_source': self.data_path
                },
                'timestamp': datetime.now().isoformat(),
                'total_problems': len(results),
                'results': results
            }
            
            with open(self.output_path, 'w', encoding='utf-8') as f:
                json.dump(output_data, f, ensure_ascii=False, indent=2)
            
            print(f"\n结果已保存到: {self.output_path}")
            
        except Exception as e:
            print(f"保存结果失败: {str(e)}")

